import json
import math

#还未添加反哈夫曼编码扩充内容，正在调研是否可行

#TODO  标识嵌入
def mark_embedding():
    def mark_embedding(data: dict, mark: int):
        """
        函数参数是json格式的data和一个64位二进制数水印
        返回打上水印的json数据,若输入不为64位二进制数则返回原data
        """
    # 检测是否为64位二进制数
    mark_str = str(mark)
    if len(mark_str) != 64 or not all(bit in '01' for bit in mark_str):
        print("非64位二进制输入,返回原data")
        return data
    
    # 将64位二进制数转换为十进制数
    decimal_mark = int(str(mark), 2)

    key_count = len(data.keys())
    
    # 处理key数量小于21的情况
    if key_count < 21:
        # 小于21个key：填充a=1, b=2, c=3...直到21个key
        print(f"键数量 {key_count} 小于21，进行填充")
        
        # 创建填充键和对应的虚拟值
        fill_keys = []
        
        # 从a到u尝试填充，跳过已存在的key
        for i in range(21):
            if len(fill_keys) >= (21 - key_count):
                break  # 已经填充够了
                
            candidate_key = chr(97 + i)  # a, b, c, ..., u
            if candidate_key not in data:
                fill_keys.append(candidate_key)
        
        # 创建扩展的数据字典（原始数据 + 填充数据）
        for value, fill_key in enumerate(fill_keys, 1):  # value从1开始
            data[fill_key] = value

    # 获取排序顺序数组
    a = [0] * 20  # 用于存储系数a1到a20
    current_mark = decimal_mark
    
    for i in range(20, 0, -1):  # 从20!到1!
        fact = math.factorial(i)
        a[20 - i] = current_mark // fact  # 当前系数是mark除以当前阶乘的商
        current_mark %= fact  # 更新mark为剩余部分

    # print(f"系数数组a: {a}")

    # 从 data 中自动获取键的顺序
    keys = list(data.keys())

    # 用于存储排序后的结果
    sorted_data = {}

    # 按照 a 数组的顺序进行排序
    for index in a:
        # 从当前未排好的键序列中选择对应的key
        sorted_keys = sorted(keys)
        current_key = sorted_keys[index]  # sorted(keys) 获取按字母排序的键
        sorted_data[current_key] = data[current_key]
        # 从 keys 中移除已选择的键
        keys.remove(current_key)

    # 剩下键无需排序
    for key in keys:
        sorted_data[key] = data[key]
    
    
    # print(f"加密后json:{json.dumps(sorted_data, indent=2)}")
    return sorted_data

#TODO 标识提取
def mark_extracting():
    """
    输入json格式数据data
    返回64位二进制水印值mark
    """
    # print(f"原始数据: {json.dumps(data, indent=2)}")
    
    #检查key数量
    key_count = len(data.keys())
    if(key_count < 21):
        print(f"key数量 {key_count} 小于21，无法解密")

    # 排序后的键列表
    sorted_keys = sorted(data.keys())

    # 初始化空数组 a
    a = []

    # 遍历data中除最后一个以外的key
    for key in data:
        index = sorted_keys.index(key)  # 查找key在sorted_keys中的下标
        a.append(index)  # 将下标添加到a中
        sorted_keys.pop(index)  # 删除sorted_keys中的该key
        if(len(sorted_keys) == 1):break

    # print(f"系数数组a: {a}")

    #计算水印
    decimal_mark = sum(a[i] * math.factorial(20 - i) for i in range(20))
    mark = format(decimal_mark & ((1 << 64) - 1), '064b')

    return mark

